#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import time

sys.path.append('.')

"""
    手工过滤不符合要求的图片的可视化工具，可同时显示同一个人的所有图片
"""


h_image_count = 10
v_image_count = 10
img_width = 96
img_height = 96

batch_size = 32


ix, iy = -1, -1
draw_image = None
image_bboxes = []
select_image_bbox_idxes = set()


def get_xy_bbox_idx(x, y):
    global image_bboxes
    bbox_idx = None
    for i, bbox in enumerate(image_bboxes):
        # print(bbox, x, y)
        if (bbox[0] < x < bbox[2]) and (bbox[1] < y < bbox[3]):
            bbox_idx = i
            break
    return bbox_idx


def mouse_event(event, x, y, flags, param):
    global draw_image, image_bboxes, select_image_bbox_idxes

    if draw_image is not None:
        if event == cv2.EVENT_LBUTTONDOWN:
            # print(x, y)
            bbox_idx = get_xy_bbox_idx(x, y)
            if bbox_idx is not None:
                if bbox_idx not in select_image_bbox_idxes:
                    select_image_bbox_idxes.add(bbox_idx)
                else:
                    select_image_bbox_idxes.remove(bbox_idx)


def show_images(image_dir, file_names, dis_text='', win_name='image', wait=False):
    h_image_count = 20
    v_image_count = 10
    img_width = 96
    img_height = 96

    display_image_count = h_image_count * v_image_count
    files = []
    total_image_count = len(file_names)
    last_idx = total_image_count - 1
    for idx, file_name in enumerate(file_names):
        files.append(file_name)
        file_count = len(files)
        if file_count < display_image_count and (idx < last_idx):
            continue
        imgs = []
        for file_name in files:
            file_path = os.path.join(image_dir, file_name.strip())
            img = cv2.imread(file_path)
            img = cv2.resize(img, (img_width, img_height))
            # w, h = img.shape[1], img.shape[0]
            # wscale = img_width / w
            # hscale = img_height / h
            # scale = min(wscale, hscale)
            # w = int(round(w * scale))
            # h = int(round(h * scale))
            # image = cv2.resize(img, (w, h))
            # img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
            # x = (img_width-w)//2
            # y = (img_height-h)//2
            # img[y:y+h, x:x+w, :] = image
            imgs.append(img)
        files.clear()
        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        draw_image = np.vstack(images)
        cv2.putText(draw_image, '%s %d/%d/%d' % (dis_text, idx + 1, last_idx + 1, total_image_count), (8, 24), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 1)
        cv2.imshow(win_name, draw_image)
        if wait:
            key = cv2.waitKey()
            if key == 27 or key == ord('q') or key == ord('Q'):
                return
        else:
            break


def read_file_list(file_list, use_select=True):
    faceid_subfaceid_file_selects = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            file_name = lines[2]
            if use_select:
                select = int(lines[3])
            else:
                select = 0

            faceid_subfaceid = '{} {}'.format(faceid, subfaceid)
            file_selects = faceid_subfaceid_file_selects.get(faceid_subfaceid, None)
            if file_selects is None:
                faceid_subfaceid_file_selects[faceid_subfaceid] = [[file_name], [select]]
            else:
                file_selects[0].append(file_name)
                file_selects[1].append(select)

        print('Read File Count: {}'.format(file_count))
    faceid_list = list(faceid_subfaceid_file_selects.keys())
    faceid_list.sort()
    return faceid_list, faceid_subfaceid_file_selects


def read_features(file_list, feature_file, feat_dim):
    faceid_subfaceid_file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            file_name = lines[2]
            # mode = int(lines[3])

            faceid_subfaceid = '{} {}'.format(faceid, subfaceid)
            # feature = np.fromfile(feat_file, dtype=np.float32, count=feat_dim).reshape(1, feat_dim)
            feature = feat_file.read(feat_dim*4)

            file_features = faceid_subfaceid_file_features.get(faceid_subfaceid, None)
            if file_features is None:
                faceid_subfaceid_file_features[faceid_subfaceid] = {file_name: feature}
            else:
                file_features[file_name] = feature

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return faceid_subfaceid_file_features


def save_process(select_file_list_path, select_file_list_map, cur_id_idx_file_path, cur_id_idx):
    if len(select_file_list_map) > 0:
        with open(select_file_list_path, 'w') as file:
            select_id_list = list(select_file_list_map.keys())
            select_id_list.sort()
            for select_id in select_id_list:
                select_file_names, is_selects = select_file_list_map[select_id]
                # select_file_names.sort()
                for select_file_name, is_select in zip(select_file_names, is_selects):
                    select_info = '{} {} {}\n'.format(select_id, select_file_name, is_select)
                    # print(select_info)
                    file.write(select_info)
        os.system('chmod a+wr {}'.format(select_file_list_path))

    with open(cur_id_idx_file_path, 'w') as file:
        process_info = '{}\n'.format(cur_id_idx)
        file.write(process_info)
    os.system('chmod a+wr {}'.format(cur_id_idx_file_path))


def cosine_distance(input1):
    distmat = 1. - np.dot(input1, input1.T)
    distmat = np.clip(distmat, 0.0, 1.0)
    return distmat


def _face_cluster(feats, left_idxes, same_image_distance_thr):
    distmat = cosine_distance(feats)
    dist_sum = distmat.sum(axis=1)
    idx = dist_sum.argmin()
    dist = distmat[idx]
    mask = dist < same_image_distance_thr
    same_face_idxes = left_idxes[mask]
    left_mask = np.logical_not(mask)
    left_face_idxes = left_idxes[left_mask]
    return left_idxes[idx], same_face_idxes, left_face_idxes, left_mask


def face_cluster(feats, same_image_distance_thr):
    same_face_idxes_list = []
    left_face_idxes = np.arange(0, feats.shape[0])
    typicial_idxes = set()
    while True:
        left_face_idxes_count = len(left_face_idxes)
        typicial_idx, same_face_idxes, left_face_idxes, left_mask = _face_cluster(feats, left_face_idxes, same_image_distance_thr)
        same_face_idxes_list.append(same_face_idxes.tolist())
        typicial_idxes.add(typicial_idx)
        if left_face_idxes.size < 1:
            break
        elif left_face_idxes.size == 1:
            typicial_idxes.add(int(left_face_idxes[0]))
            same_face_idxes_list.append(left_face_idxes.tolist())
            break
        elif left_face_idxes.size == left_face_idxes_count:
            for left_face_idx in left_face_idxes:
                left_face_idx = int(left_face_idx)
                typicial_idxes.add(left_face_idx)
                same_face_idxes_list.append([left_face_idx])
            break
        feats = feats[left_mask]
    return same_face_idxes_list, typicial_idxes


def face_rec(image_dir, file_list, file_features, feature_dim, same_face_score_thr, show=False, max_display_count=20):
    face_count = len(file_list)
    face_nums = [0] * face_count
    typicial_idxes = set()
    if file_features is not None:
        features = [np.frombuffer(file_features[file_name], dtype=np.float32, count=feature_dim)
                    for file_name in file_list]
        features = np.stack(features)
        same_image_distance_thr = 1.0 - same_face_score_thr
        same_face_idxes_list, typicial_idxes = face_cluster(features, same_image_distance_thr)
        if show:
            cv2.destroyAllWindows()
            cv2.namedWindow('image')
            cv2.setMouseCallback('image', mouse_event)
        for num, idxes in enumerate(same_face_idxes_list):
            files = []
            for idx in idxes:
                face_nums[idx] = num
                files.append(file_list[idx])
            if show and (num < max_display_count):
                show_images(image_dir, files, win_name='{}/{}'.format(num, face_count-1))
    return face_nums, typicial_idxes


def filter_select_face(output_dir, select_file_list_map, faceid_file_features):
    filter_file_list_path = os.path.join(output_dir, 'FilterSelectFace.txt')
    filter_features_path = os.path.join(output_dir, 'FilterSelectFeature.dat')
    filter_file_list_file = open(filter_file_list_path, 'w')
    filter_features_file = open(filter_features_path, 'wb')
    face_count = 0
    face_image_count = 0
    for faceid, [files, selects] in select_file_list_map.items():
        file_features = faceid_file_features[faceid]
        has_face_image = False
        for file_name, is_select in zip(files, selects):
            if 0 != is_select:
                continue
            out_info = '{} {} {}\n'.format(faceid, file_name, is_select)
            # print(out_info)
            filter_file_list_file.write(out_info)
            feature = file_features[file_name]
            filter_features_file.write(feature)
            face_image_count += 1
            has_face_image = True
        if has_face_image:
            face_count += 1
    filter_file_list_file.close()
    filter_features_file.close()
    os.system('chmod a+wr {}'.format(filter_file_list_path))
    os.system('chmod a+wr {}'.format(filter_features_path))
    print('Face Count: {}, Face Image Count: {}'.format(face_count, face_image_count))


def read_faceid_list(faceid_list):
    faceids_set = set()
    with open(faceid_list, 'r') as file:
        for line in file.readlines():
            faceid = line.strip()
            faceids_set.add(faceid)
    return faceids_set


def main():
    global draw_image, image_bboxes, select_image_bbox_idxes
    parser = argparse.ArgumentParser(description="Select Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data1/data/VideoFaceImage", help="path to image file", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/Feature.select.FR.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/Feature.select.FR.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/VideoFaceImage", help="path to mot reid image info", type=str
    )
    parser.add_argument(
        "--faceid_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/FaceIdList.txt", help="path to select face id list", type=str
    )
    parser.add_argument("--use_temp", action='store_true', help="use temp save, if crash can use it")
    parser.add_argument('--save_time', type=int, default=60, help='save time, second ')
    parser.add_argument("--use_feature", action='store_true', help="use feature")
    parser.add_argument("--feature_dim", default=512, help="feature dim", type=int)
    parser.add_argument('--same_face_score_thr', type=float, default=0.6, help='same face score thr')
    parser.add_argument("--show", action='store_true', help="show face rec")
    parser.add_argument("--show_count", default=20, help="show face rec id max count", type=int)
    parser.add_argument("--filter", action='store_true', help="filter select face")
    parser.add_argument("--use_faceid_list", action='store_true', help="use faceid list")
    parser.add_argument("--select_count", default=2, help="select min face image count", type=int)
    args = parser.parse_args()

    save_time = args.save_time
    image_dir = args.image_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    display_image_count = h_image_count * v_image_count

    select_file_list_path = os.path.join(output_dir, 'select_reid_images.txt')
    if not os.path.exists(select_file_list_path):
        id_list, select_file_list_map = read_file_list(args.file_list, False)
    else:
        id_list, select_file_list_map = read_file_list(select_file_list_path)

    same_face_score_thr = args.same_face_score_thr
    if args.use_feature or args.filter:
        faceid_file_features = read_features(args.file_list, args.feature_file, args.feature_dim)

    if args.filter:
        filter_select_face(output_dir, select_file_list_map, faceid_file_features)
        return

    if args.use_faceid_list:
        faceids_set = read_faceid_list(args.faceid_list)
        process_faceid_count = 0

    cv2.namedWindow('image')
    cv2.setMouseCallback('image', mouse_event)

    id_count = len(id_list)
    last_id = id_list[-1]
    last_id_idx = id_count - 1

    cur_id_idx = 0
    cur_id_idx_file_path = os.path.join(output_dir, 'cur_id_idx.txt')
    if os.path.exists(cur_id_idx_file_path):
        with open(cur_id_idx_file_path, 'r') as file:
            for line in file.readlines():
                cur_id_idx = int(line.strip())
                break

    if args.use_temp:
        tmp_output_dir = os.path.join('/dev/shm', 'SelectImage')
        if not os.path.exists(tmp_output_dir):
            os.makedirs(tmp_output_dir)
            os.system('chmod a+wr {}'.format(tmp_output_dir))

        tmp_select_file_list_path = os.path.join(tmp_output_dir, 'select_reid_images.txt')
        tmp_cur_id_idx_file_path = os.path.join(tmp_output_dir, 'cur_id_idx.txt')

    select_change = False
    start_save_time = time.time()
    while cur_id_idx <= last_id_idx:
        cur_trk_id = id_list[cur_id_idx]
        if args.use_faceid_list:
            if process_faceid_count >= len(faceids_set):
                cur_id_idx = last_id_idx
                break
            if cur_trk_id in faceids_set:
                process_faceid_count += 1
            else:
                cur_id_idx += 1
                continue
        print('{}/{}'.format(cur_trk_id, last_id))
        file_list, select_list = select_file_list_map[cur_trk_id]
        if len(file_list) < args.select_count:
            cur_id_idx += 1
            continue
        face_nums = [0] * len(file_list)
        typicial_face_idxes = set()
        if args.use_feature:
            file_features = faceid_file_features.get(cur_trk_id, None)
            face_nums, typicial_face_idxes = face_rec(image_dir, file_list, file_features, args.feature_dim, same_face_score_thr, False, 0)
        last_idx = len(file_list) - 1
        start_idx = 0
        files = []

        for idx, file_name in enumerate(file_list):
            files.append(file_name)
            file_count = len(files)
            if file_count == 1:
                start_idx = idx
            if file_count < display_image_count and (idx < last_idx):
                continue
            imgs = []
            for file_name in files:
                file_path = os.path.join(image_dir, file_name)
                img = cv2.imread(file_path)
                img = cv2.resize(img, (img_width, img_height))
                # w, h = img.shape[1], img.shape[0]
                # wscale = img_width / w
                # hscale = img_height / h
                # scale = min(wscale, hscale)
                # w = int(round(w * scale))
                # h = int(round(h * scale))
                # image = cv2.resize(img, (w, h))
                # img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
                # x = (img_width-w)//2
                # y = (img_height-h)//2
                # img[y:y+h, x:x+w, :] = image
                imgs.append(img)

            image_count = len(imgs)
            images = []
            x_off = 0
            y_off = 0
            image_bboxes.clear()
            select_image_bbox_idxes.clear()
            for select_idx in range(0, image_count):
                if 0 != select_list[start_idx+select_idx]:
                    select_image_bbox_idxes.add(select_idx)
            for j, img in enumerate(imgs):
                if j % h_image_count == 0:
                    x_off = 0
                    if j != 0:
                        y_off += img_height
                    h_images = [img]
                else:
                    x_off += img_width
                    h_images.append(img)
                image_bboxes.append([x_off, y_off, (x_off + img_width), (y_off + img_height)])
                if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                    images.append(np.concatenate(h_images, axis=1))

            if len(images) > 1:
                w0 = images[0].shape[1]
                h0 = images[0].shape[0]
                w1 = images[-1].shape[1]
                h1 = images[-1].shape[0]
                if (w0 != w1) or (h0 != h1):
                    # image = np.empty_like(images[0])
                    image = np.zeros_like(images[0])
                    image[0:h1, 0:w1, :] = images[-1]
                    images[-1] = image
            draw_image = np.vstack(images)

            change_trk_id = False
            is_exit = False
            thr_str = ''
            while True:
                image = draw_image.copy()
                for select_idx in select_image_bbox_idxes:
                    bbox = image_bboxes[select_idx]
                    cv2.circle(image, (bbox[0]+14, bbox[1]+14), 8, (0, 0, 255), -1)
                if args.use_feature:
                    for select_idx in range(0, image_count):
                        bbox = image_bboxes[select_idx]
                        face_idx = start_idx + select_idx
                        face_num = face_nums[face_idx]  # faceID
                        if face_idx in typicial_face_idxes:
                            cv2.circle(image, (bbox[2] - 10, bbox[3] - 10), 8, (0, 255, 0), -1)
                        if face_num != 0:   # 不显示等于0的，方便查看差异大的人脸图片，一般数字越大差异越大
                            cv2.putText(image, str(face_num), (bbox[0]+2, bbox[3]-4), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), thickness=2)
                cv2.putText(image, '%d/%d %d/%d %s %s' % (cur_id_idx+1, last_id_idx+1, idx + 1, last_idx + 1, same_face_score_thr, thr_str),
                            (8, 24), cv2.FONT_HERSHEY_PLAIN, 2.0, (0, 255, 0), thickness=2)
                cv2.imshow('image', image)
                key = cv2.waitKey(50)
                if (key == 27) or (key == ord('q')) or (key == ord('Q')):
                    is_exit = True
                    break
                elif key == ord(' '):  # 同一ID的下一页
                    select_change = True
                    break
                elif key == ord('a') or key == ord('A'):  # 上一个ID
                    select_change = True
                    if cur_id_idx > 0:
                        cur_id_idx -= 1
                        change_trk_id = True
                    break
                elif key == ord('d') or key == ord('D'):  # 下一个ID
                    select_change = True
                    change_trk_id = True
                    cur_id_idx += 1
                    break
                elif key == ord('r') or key == ord('R'):  # 选中当前页典型人脸
                    for select_idx in range(0, image_count):
                        face_idx = start_idx + select_idx
                        if face_idx in typicial_face_idxes:
                            select_image_bbox_idxes.add(select_idx)
                    continue
                elif key == ord('e') or key == ord('E'):  # 当前页全选
                    select_image_bbox_idxes.update(range(0, image_count))
                    continue
                elif key == ord('w') or key == ord('W'):  # 当前页全不选
                    select_image_bbox_idxes.clear()
                    continue
                elif (ord('0') <= key <= ord('9')) or (key == ord('.')):  # 输入新阈值
                    if thr_str == '':
                        thr_str = chr(key)
                    else:
                        thr_str += chr(key)
                    continue
                elif key == ord('f') or key == ord('F'):  # 按照设定的阈值对人脸做重新识别
                    if thr_str != '':
                        same_face_score_thr = float(thr_str)
                        if same_face_score_thr < 0.0:
                            same_face_score_thr = 0.0
                        elif same_face_score_thr > 1.0:
                            same_face_score_thr = 1.0
                        thr_str = ''
                    if args.use_feature:
                        file_features = faceid_file_features.get(cur_trk_id, None)
                        face_nums, typicial_face_idxes = face_rec(image_dir, file_list, file_features, args.feature_dim, same_face_score_thr, args.show, args.show_count)
                    continue
                elif key == ord('c') or key == ord('C'):
                    cv2.destroyAllWindows()
                    cv2.namedWindow('image')
                    cv2.setMouseCallback('image', mouse_event)
                    continue
                elif key == ord('s') or key == ord('S'):  # Save
                    if args.use_temp:
                        save_process(tmp_select_file_list_path, select_file_list_map, tmp_cur_id_idx_file_path, cur_id_idx)
                    start_save_time = time.time()
                    select_change = False
                    break
                if args.use_temp:
                    # 定时保存
                    if time.time() - start_save_time >= save_time:
                        if select_change:
                            save_process(tmp_select_file_list_path, select_file_list_map, tmp_cur_id_idx_file_path, cur_id_idx)
                            select_change = False
                        start_save_time = time.time()
            # if len(select_image_bbox_idxes) > 0:
            for bbox_idx in range(0, image_count):
                select_idx = start_idx + bbox_idx
                if bbox_idx in select_image_bbox_idxes:
                    select_list[select_idx] = 1
                else:
                    select_list[select_idx] = 0
            files.clear()
            if is_exit:
                save_process(select_file_list_path, select_file_list_map, cur_id_idx_file_path, cur_id_idx)
                return
            if change_trk_id:
                break
    save_process(select_file_list_path, select_file_list_map, cur_id_idx_file_path, cur_id_idx)
    print('finish!')


if __name__ == '__main__':
    main()
